import numpy as np

class rank_ensemble(object):

    '''
    usage:
    model = rank_ensemble
    model_para = {
        'models': [m1, m2],
        'weight': [10, 9]
    }
    '''

    def __init__(self, models, weight=None):
        self.models = models
        if weight is None:
            self.weight = np.zeros((len(self.models), ))
            weight[:] = 1 / len(self.models)
        else:
            self.weight = np.array(weight)
            self.weight = self.weight / self.weight.sum()

    def fit(self, X_train, y_train):
        for model in self.models:
            model.fit(X_train, y_train)

    def predict_rank(self, X_test):
        normalized_score = np.zeros((X_test.shape[0], ))
        for model_index, model in enumerate(self.models):
            y_pred = model.predict_proba(X_test)
            y_pred = y_pred[:, 1]
            pred_rank = np.argsort(y_pred)
            for i, index in enumerate(pred_rank):
                normalized_score[index] += i / len(pred_rank) * self.weight[model_index]
        return normalized_score / len(self.models)

    def predict_proba(self, X_test):
        result = np.zeros((X_test.shape[0], 2))
        result[:, 1] = self.predict_rank(X_test)
        result[:, 0] = 1 - result[:, 1]
        return result

